package org.finesys.ai.controller.chat;

import io.swagger.v3.oas.annotations.security.SecurityRequirement;
import lombok.RequiredArgsConstructor;
import org.finesys.ai.constants.AigcControllerRestApiConfig;
import org.finesys.ai.constants.EmbedConst;
import org.finesys.ai.entity.AigcDoc;
import org.finesys.ai.service.AigcDocService;
import org.finesys.ai.service.EmbeddingService;
import org.finesys.common.core.exception.BusinessException;
import org.finesys.common.core.module.R;
import org.finesys.common.core.support.concurrent.TaskManager;
import org.finesys.common.file.FileInfo;
import org.finesys.common.file.client.FileTemplate;
import org.finesys.common.file.properties.FileProperties;
import org.finesys.common.security.core.util.SecurityUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.util.concurrent.Executors;


@RequestMapping(value = AigcControllerRestApiConfig.EMBEDDING)
@RestController
@RequiredArgsConstructor
@SecurityRequirement(name = HttpHeaders.AUTHORIZATION)
public class EmbeddingController {

    private final EmbeddingService embeddingService;
    private final AigcDocService aigcDocService;
    private final FileTemplate fileTemplate;
    private final FileProperties fileProperties;

    @PostMapping("/text")
    @PreAuthorize("@pms.hasPermission('aigc_embedding_text')")
    public R<?> text(AigcDoc aigcDoc) {
        if (!StringUtils.hasText(aigcDoc.getContent())) {
            return R.failed("内容不能为空");
        }
        aigcDoc.setType(EmbedConst.INPUT);
        aigcDocService.save(aigcDoc);
        try {
            embeddingService.embedDocsSlice(aigcDoc);
        } catch (Exception e) {
            aigcDocService.removeById(aigcDoc.getId());
        }
        return R.success();
    }

    @PostMapping("/doc/{knowledgeId}")
    @PreAuthorize("@pms.hasPermission('aigc_embedding_file')")
    public R<?> embedDoc(MultipartFile multipartFile, @PathVariable String knowledgeId) throws Exception {
        //获取当前用户ID
        String userId = SecurityUtils.getUserId();
        // 文件上传
        FileInfo fileInfo = fileTemplate.putObject(fileProperties.getBucketName(), multipartFile.getName(), multipartFile.getInputStream());
        //存储文件信息到数据库
        AigcDoc aigcDoc = new AigcDoc()
                .setType(EmbedConst.UPLOAD)
                .setName(multipartFile.getName())
                .setOrigin(multipartFile.getOriginalFilename())
                .setKnowledgeId(knowledgeId)
                .setFileSize(fileInfo.getSize())
                .setUrl(fileInfo.getUrl());
        aigcDocService.save(aigcDoc);
        // 嵌入向量
        TaskManager.submitTask(userId, Executors.callable(() -> embeddingService.embedDocsSlice(aigcDoc)));
        return R.success();
    }

    @GetMapping("/reEmbed/{docsId}")
    @PreAuthorize("@pms.hasPermission('aigc_embedding_reEmbed')")
    public R<?> reEmbed(@PathVariable String docsId) {
        String userId = SecurityUtils.getUserId();
        AigcDoc docs = aigcDocService.getById(docsId);
        if (docs == null) {
            throw new BusinessException("没有查询到文档数据");
        }
        if (EmbedConst.INPUT.equals(docs.getType())) {
            text(docs);
        }
        if (EmbedConst.UPLOAD.equals(docs.getType())) {
            // clear before re-embed
            embeddingService.clearDocSlices(docsId);
            TaskManager.submitTask(userId, Executors.callable(() -> {
                embeddingService.embedDocsSlice(docs);
            }));
        }
        return R.success();
    }

    @PostMapping("/search")
    public R<?> search(@RequestBody AigcDoc data) {
        return R.success(embeddingService.search(data));
    }
}